package trianing.taxi;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.google.common.hash.BloomFilter;
import com.google.common.hash.Funnels;
import org.apache.flink.api.common.functions.FilterFunction;
import org.apache.flink.api.common.functions.RuntimeContext;
import org.apache.flink.api.common.serialization.SimpleStringSchema;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.runtime.state.hashmap.HashMapStateBackend;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.KeyedProcessFunction;
import org.apache.flink.streaming.connectors.elasticsearch.ElasticsearchSinkFunction;
import org.apache.flink.streaming.connectors.elasticsearch.RequestIndexer;
import org.apache.flink.streaming.connectors.elasticsearch7.ElasticsearchSink;
import org.apache.flink.streaming.connectors.kafka.FlinkKafkaConsumer;
import org.apache.flink.util.Collector;
import org.apache.http.HttpHost;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.client.Requests;

import java.nio.charset.StandardCharsets;
import java.util.*;

/**
 * @desc
 * @Author xuyouchang
 * @Date 2022/2/16
 */
public class EtlByStreamApi {

    private static class DeDuplicateKeyedProcessFun extends KeyedProcessFunction<Long, String, String> {
        // 布隆过滤器期望最大数据量
        private static final int BF_CARDINAL_THRESHOLD = 1000000;
        // 假阳性比例
        private static final double BF_FALSE_POSITIVE_RATE = 0.01;

        private volatile BloomFilter<byte[]> subOrderFilter;

        // 根据当前时间戳获取第二天0时0分0秒的时间戳
        public long tomorrowZeroTimestampMs(long now, int timeZone) {
            return now - (now + timeZone * 3600000) % 86400000 + 86400000;
        }

        @Override
        public void open(Configuration parameters) throws Exception {
            subOrderFilter = BloomFilter.create(Funnels.byteArrayFunnel(), BF_CARDINAL_THRESHOLD, BF_FALSE_POSITIVE_RATE);
        }

        @Override
        public void onTimer(long timestamp, OnTimerContext ctx, Collector<String> out) throws Exception {
            subOrderFilter = BloomFilter.create(Funnels.byteArrayFunnel(), BF_CARDINAL_THRESHOLD, BF_FALSE_POSITIVE_RATE);
        }

        @Override
        public void processElement(String value, Context ctx, Collector<String> out) throws Exception {
            JSONObject valueObject = JSON.parseObject(value);
            String id = valueObject.getString("rideId") + "-"
                    + valueObject.getString("taxiId") + "-"
                    + valueObject.getString("driverId");
            if (!subOrderFilter.mightContain(id.getBytes(StandardCharsets.UTF_8))) {
                subOrderFilter.put(id.getBytes(StandardCharsets.UTF_8));
                out.collect(value);
            }
            ctx.timerService().registerProcessingTimeTimer(tomorrowZeroTimestampMs(System.currentTimeMillis(), 8) + 1);
        }

    }

    public static void main(String[] args) throws Exception {
        Properties properties = new Properties();
        properties.setProperty("bootstrap.servers", "172.17.1.101:9092");
        properties.setProperty("group.id", "test13");
        properties.setProperty("auto.offset.reset", "earliest");

        Configuration configuration = new Configuration();
        StreamExecutionEnvironment streamExecutionEnvironment = StreamExecutionEnvironment.createLocalEnvironmentWithWebUI(configuration);
        streamExecutionEnvironment.enableCheckpointing(5000L);
        streamExecutionEnvironment.setStateBackend(new HashMapStateBackend());
        DataStreamSource<String> dataStream = streamExecutionEnvironment.addSource(new FlinkKafkaConsumer<>("taxi_ride",
                new SimpleStringSchema(), properties));

        DataStream<String> dataStream1 = dataStream
                .filter(new FilterFunction<String>() {
                    @Override
                    public boolean filter(String data) throws Exception {
                        JSONObject jsonObject = JSONObject.parseObject(data);
                        return jsonObject.getLong("taxiId") > 0 &&
                                jsonObject.getLong("driverId") > 0 &&
                                jsonObject.getLong("rideId") > 0 &&
                                jsonObject.getShort("passengerCnt") > 0;
                    }
                }).keyBy(value -> JSONObject.parseObject(value).getLong("taxiId"))
                .process(new DeDuplicateKeyedProcessFun());

        List<HttpHost> httpHosts = new ArrayList<>();
        httpHosts.add(new HttpHost("172.17.1.101", 9200, "http"));
        httpHosts.add(new HttpHost("172.17.1.102", 9200, "http"));
        httpHosts.add(new HttpHost("172.17.1.103", 9200, "http"));

        // 使用 ElasticsearchSink.Builder 创建 ElasticsearchSink
        ElasticsearchSink.Builder<String> esSinkBuilder = new ElasticsearchSink.Builder<>(
                httpHosts,
                new ElasticsearchSinkFunction<String>() {
                    public IndexRequest createIndexRequest(String element) {
                        JSONObject value = JSONObject.parseObject(element);
                        return Requests.indexRequest()
                                .index("dd_taxi_ride_01")
                                .id(value.getString("taxiId").concat(value.getString("driverId")).concat(value.getString("rideId")))
                                .source(value);
                    }

                    @Override
                    public void process(String element, RuntimeContext ctx, RequestIndexer indexer) {
                        indexer.add(createIndexRequest(element));
                    }
                }
        );
        dataStream1.addSink(esSinkBuilder.build()).uid("写入es");
        streamExecutionEnvironment.execute("数据过滤");
    }
}
